# -*- coding: utf-8 -*-
"""
问答数据存储模块

本模块负责管理问答系统的数据存储，提供标准的问答数据集和数据管理接口。
主要用于存储、检索和管理问答对，支持基于关键词的快速查找。

核心功能：
1. 多知识库支持：支持从多个JSON文件加载问答数据
2. 数据管理：支持增删改查操作
3. 关键词索引：基于关键词的快速检索
4. 数据持久化：支持数据的保存和加载（可扩展）

数据结构设计：
每个问答项包含以下字段：
- question (str): 问题文本
- answer (str): 答案文本  
- key (str): 唯一标识符
- keywords (List[str]): 关键词列表，用于检索匹配

适用场景：
- 问答系统的基础数据源
- 知识库构建
- 测试数据管理
- 示例数据提供
"""
import json
import os
from pathlib import Path
from typing import Dict, List, Optional, Union
from loguru import logger


class QAStorage:
    """
    问答数据存储管理器 - 支持多知识库文件加载
    """

    def __init__(self, knowledge_bases: Optional[List[str]] = None, data_dir: str = "data/raw"):
        """
        初始化问答存储管理器
        
        Args:
            knowledge_bases: 知识库文件名列表，如 ["logic_qa.json", "flood_qa.json"]
            data_dir: 数据文件目录
        """
        self.data_dir = Path(data_dir)
        self.data = []
        self.knowledge_bases = knowledge_bases or []
        
        if self.knowledge_bases:
            self._load_knowledge_bases()
        else:
            logger.info("QAStorage初始化完成，未加载任何知识库文件")
    
    def _load_knowledge_bases(self):
        """从文件加载知识库数据"""
        total_loaded = 0
        
        for kb_file in self.knowledge_bases:
            file_path = self.data_dir / kb_file
            
            if not file_path.exists():
                logger.warning(f"知识库文件不存在: {file_path}")
                continue
                
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    kb_data = json.load(f)
                    
                # 验证数据格式
                valid_data = self._validate_qa_data(kb_data, kb_file)
                self.data.extend(valid_data)
                total_loaded += len(valid_data)
                
                logger.info(f"成功加载知识库 {kb_file}，包含 {len(valid_data)} 条数据")
                
            except Exception as e:
                logger.error(f"加载知识库文件 {kb_file} 失败: {e}")
        
        logger.info(f"知识库加载完成，总共加载了 {total_loaded} 条问答数据")
    
    def _validate_qa_data(self, data: List[Dict], source_file: str) -> List[Dict]:
        """验证问答数据格式"""
        valid_data = []
        required_fields = ['question', 'answer', 'key', 'keywords']
        
        for i, item in enumerate(data):
            if not isinstance(item, dict):
                logger.warning(f"{source_file} 第 {i+1} 条数据格式错误：不是字典类型")
                continue
                
            # 检查必需字段
            missing_fields = [field for field in required_fields if field not in item]
            if missing_fields:
                logger.warning(f"{source_file} 第 {i+1} 条数据缺少字段: {missing_fields}")
                continue
            
            # 添加默认字段
            item.setdefault('category', 'default')
            item.setdefault('source', source_file)
            
            valid_data.append(item)
        
        return valid_data
    
    def add_knowledge_base(self, kb_file: str) -> bool:
        """动态添加知识库文件"""
        if kb_file in self.knowledge_bases:
            logger.warning(f"知识库 {kb_file} 已存在")
            return False
            
        file_path = self.data_dir / kb_file
        if not file_path.exists():
            logger.error(f"知识库文件不存在: {file_path}")
            return False
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                kb_data = json.load(f)
            
            valid_data = self._validate_qa_data(kb_data, kb_file)
            self.data.extend(valid_data)
            self.knowledge_bases.append(kb_file)
            
            logger.info(f"成功添加知识库 {kb_file}，包含 {len(valid_data)} 条数据")
            return True
            
        except Exception as e:
            logger.error(f"添加知识库 {kb_file} 失败: {e}")
            return False
    
    def get_knowledge_bases_info(self) -> Dict[str, Union[List[str], int, Dict]]:
        """获取知识库信息"""
        category_stats = {}
        source_stats = {}
        
        for item in self.data:
            category = item.get('category', 'default')
            source = item.get('source', 'unknown')
            
            category_stats[category] = category_stats.get(category, 0) + 1
            source_stats[source] = source_stats.get(source, 0) + 1
        
        return {
            "loaded_knowledge_bases": self.knowledge_bases,
            "total_qa_count": len(self.data),
            "category_distribution": category_stats,
            "source_distribution": source_stats
        }
    
    def get_all_qa(self) -> List[Dict]:
        """获取所有问答数据"""
        logger.debug(f"获取所有问答数据，共 {len(self.data)} 条")
        return self.data.copy()

    def add_qa(self, question: str, answer: str, key: str, keywords: List[str], 
               category: str = None, source: str = None) -> bool:
        """添加新的问答数据"""
        # 参数验证
        if not question or not question.strip():
            raise ValueError("问题不能为空")
        if not answer or not answer.strip():
            raise ValueError("答案不能为空")
        if not key or not key.strip():
            raise ValueError("key不能为空")
        if not isinstance(keywords, list):
            raise ValueError("关键词必须是列表格式")

        # 检查key唯一性
        if self.get_qa_by_key(key) is not None:
            raise ValueError(f"key '{key}' 已存在")

        qa_item = {
            "question": question.strip(),
            "answer": answer.strip(),
            "key": key.strip(),
            "keywords": [kw.strip() for kw in keywords if kw.strip()],
            "category": category or 'manual',
            "source": source or 'runtime'
        }

        self.data.append(qa_item)
        logger.info(f"添加问答数据成功，key: {key}")
        return True

    def get_qa_by_key(self, key: str) -> Optional[Dict]:
        """根据key获取问答数据"""
        if not key:
            return None

        for item in self.data:
            if item['key'] == key:
                logger.debug(f"找到key为 {key} 的问答数据")
                return item.copy()

        logger.debug(f"未找到key为 {key} 的问答数据")
        return None

    def search_by_keywords(self, keywords: Union[str, List[str]]) -> List[Dict]:
        """根据关键词搜索问答数据"""
        if isinstance(keywords, str):
            keywords = [keywords]

        if not keywords:
            return []

        results = []
        for item in self.data:
            # 检查问答记录的关键词是否包含搜索关键词
            for search_kw in keywords:
                if any(search_kw in qa_kw for qa_kw in item['keywords']):
                    results.append(item.copy())
                    break
                # 也在问题和答案中搜索
                if (search_kw in item['question'] or search_kw in item['answer']):
                    results.append(item.copy())
                    break

        logger.debug(f"关键词搜索完成，找到 {len(results)} 条记录")
        return results

    def remove_qa(self, key: str) -> bool:
        """删除指定的问答数据"""
        for i, item in enumerate(self.data):
            if item['key'] == key:
                removed_item = self.data.pop(i)
                logger.info(f"删除问答数据成功，key: {key}")
                return True

        logger.warning(f"未找到要删除的问答数据，key: {key}")
        return False

    def get_statistics(self) -> Dict[str, Union[int, List[str]]]:
        """获取存储统计信息"""
        all_keywords = []
        for item in self.data:
            all_keywords.extend(item['keywords'])

        return {
            "total_qa_count": len(self.data),
            "total_keywords": len(set(all_keywords)),
            "unique_keywords": list(set(all_keywords)),
            "average_keywords_per_qa": len(all_keywords) / len(self.data) if self.data else 0
        }
